import re
import logging
from typing import Union
from langchain.agents import AgentOutputParser
from langchain.schema import AgentAction, AgentFinish

logger = logging.getLogger(__name__)

class CustomOutputParser(AgentOutputParser):
    """自定义输出解析器，处理智谱AI的输出格式"""
    
    def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
        logger.info(f"解析Agent输出: {text}")
        
        # 检查是否包含Final Answer
        if "Final Answer:" in text:
            # 提取最终答案
            final_answer_match = re.search(r"Final Answer:\s*(.*)", text, re.DOTALL)
            if final_answer_match:
                answer = final_answer_match.group(1).strip()
                logger.info(f"提取到最终答案: {answer}")
                return AgentFinish(
                    return_values={"output": answer},
                    log=text,
                )
        
        # 检查是否包含Action和Action Input
        action_match = re.search(r"Action:\s*(.*?)\nAction Input:\s*(.*)", text, re.DOTALL)
        if action_match:
            action = action_match.group(1).strip()
            action_input = action_match.group(2).strip().strip('"')
            logger.info(f"提取到Action: {action}, Input: {action_input}")
            return AgentAction(tool=action, tool_input=action_input, log=text)
        
        # 如果无法解析，返回最终答案
        logger.warning(f"无法解析Agent输出，直接返回文本: {text}")
        return AgentFinish(
            return_values={"output": text},
            log=text,
        )